#!/bin/bash

voxceleb1_path=/ceph/home/zhangy20/datasets/VoxCeleb/voxceleb1
voxceleb2_path=/ceph/home/zhangy20/datasets/VoxCeleb/voxceleb2
musan_path=~/datasets/musan
rirs_path=~/datasets/RIRS_NOISES

trials_path=data/trials.lst

nnet_type=ResNetSE34L
loss_type=amsoftmax

. ./path.sh

stage=3
echo stage $stage

# format data dir structure by soft link
if [ $stage -eq 0 ];then
    if [ ! -d data/wav_files ]; then
        mkdir -p data/wav_files
    fi

    rm -rf data/wav_files/dev
    mkdir -p data/wav_files/dev
    # format voxceleb1
    ln -s ${voxceleb1_path}/vox1_dev_wav/* data/wav_files/dev/

    # format voxceleb2
    ln -s ${voxceleb2_path}/dev/aac/* data/wav_files/dev/

    rm -rf voxceleb1_test_v2.txt
    wget https://openslr.magicdatatech.com/resources/49/voxceleb1_test_v2.txt
    mv voxceleb1_test_v2.txt data/voxceleb1_test_v2.txt
fi


# build data list
if [ $stage -eq 1 ];then
    echo build dev data list
    python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir data/wav_files/dev \
        --data_list_path data/dev_list.csv

    echo build musan data list
    python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
        --extension wav \
        --dataset_dir $musan_path \
        --data_list_path data/musan_list.csv

    echo build rirs data list
    python3 $SPEAKER_TRAINER_ROOT/scripts/build_datalist.py \
	    --extension wav \
	    --dataset_dir $rirs_path \
	    --data_list_path data/rirs_list.csv

    rm -rf $trials_path
    python3 local/format_trials.py \
	    --voxceleb1_root $voxceleb1_path \
	    --src_trials_path data/voxceleb1_test_v2.txt \
	    --dst_trials_path $trials_path
fi


if [ $stage -eq 2 ];then
    CUDA_VISIBLE_DEVICES=0 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
	    --nnet_type $nnet_type \
	    --loss_type $loss_type \
	    --batch_size 200 \
	    --num_workers 80 \
	    --save_top_k 150 \
	    --train_list_path data/dev_list.csv \
	    --musan_list_path data/musan_list.csv \
	    --rirs_list_path data/rirs_list.csv \
	    --max_epochs 150 \
	    --max_frames 201 --min_frames 200 \
	    --learning_rate 0.0001 \
	    --distributed_backend dp \
	    --max_seg_per_spk 500 \
	    --trials_path $trials_path \
	    --eval_interval -1 \
	    --nPerSpeaker 1 \
	    --reload_dataloaders_every_epoch \
	    --gpus 1
fi


if [ $stage -eq 3 ];then
    ckpt_path=/ceph/home/zhangy20/openasv/egs/VoxCeleb/v1/exp/ResNetSE34L_amsoftmaxproto/epoch=80_loss=1.47.ckpt
    ckpt_path=/ceph/home/zhangy20/openasv/egs/VoxCeleb/v1/exp/ResNetSE34L_amsoftmaxproto/epoch=99_loss=1.54.ckpt
    ckpt_path=/ceph/home/zhangy20/speaker_trainer/examples/VoxCeleb/v1/exp/attention_amsoftmax/epoch=12_loss=1.46.ckpt
    ckpt_path=/ceph/home/zhangy20/speaker_trainer/examples/VoxCeleb/v1/exp/attention_amsoftmax/epoch=50_loss=0.62.ckpt
    ckpt_path=~/openasv/egs/VoxCeleb/v1/exp/ResNetSE34L_amsoftmax/epoch=50_loss=1.15.ckpt
    CUDA_VISIBLE_DEVICES=3 python3 -W ignore $SPEAKER_TRAINER_ROOT/main.py \
        --batch_size 64 \
        --nnet_type $nnet_type \
        --num_workers 100 \
        --train_list_path data/dev_list \
        --trials_path data/trials.lst \
        --gpus 1 \
        --max_frames 401 --min_frames 400 \
        --checkpoint_path $ckpt_path \
        --evaluate
   rm -rf lightning_logs
fi

